import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import math
# torch中变量封装函数Variable.
from torch.autograd import Variable
import numpy as np

# 位置编码器层
# 词汇在不同位置，可能会对应不同语义，
# 在Transformer的编码器结构中, 没有针对词汇位置信息
# 因此需要在Embedding层后加入一个位置编码器，
# 将以上信息嵌入到input张量中。


# 位置编码器类
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout_ratio, max_len=5000):
        """
        :param d_model: 词嵌入维度
        :param dropout_ratio: 失活比例，置0比率
        :param max_len: 每个句子的最大长度
        """

        super(PositionalEncoding, self).__init__()

        # 实例化nn中预定义的Dropout层，并将dropout传入其中，获得对象self.dropout
        self.dropout = nn.Dropout(p=dropout_ratio)

        # 初始化一个位置编码矩阵，它是一个0矩阵，矩阵的大小是max_len（行） x d_model（列）
        pe = torch.zeros(max_len, d_model)

        # 初始化一个绝对位置矩阵，大小是 max_len x 1
        # 用索引表示词汇的绝对位置用
        # 所以我们首先使用arange方法获得一个连续自然数向量，然后再使用unsqueeze方法拓展向量维度
        # 又因为参数传的是1，代表矩阵拓展的位置，会使向量变成一个max_len x 1 的矩阵,
        position = torch.arange(0, max_len).unsqueeze(1)

        # 将Pe（Positional Encoding） 和 Position 进行结合
        # 也就是将绝对位置矩阵的位置信息，加入到位置编码矩阵中，
        # 最简单思路就是先将max_len x 1的绝对位置矩阵，变换成max_len x d_model形状，然后覆盖原来的初始位置编码矩阵
        # 要做这种矩阵变换，就需要一个1xd_model形状的变换矩阵div_term，

        # 定义变化矩阵 div_term，并进行“跳跃式初始化”
        # 我们对这个变换矩阵的要求除了形状外，还希望它能够将自然数的绝对位置编码缩放成足够小的数字（乘以缩放数字），有助于在之后的梯度下降过程中更快的收敛。 这样我们就可以开始初始化这个变换矩阵了
        # 首先使用arange获得一个自然数矩阵，但是细心的同学们会发现，我们这里并没有按照预计的一样初始化一个 1 X dmodel 的矩阵，
        # 而是有了一个跳跃（下文中的2，512变256），只初始化了一半即1 x d_mode1/2 的矩阵。为什么是一半呢，其实这里并不是真正意义上的初始化了一半的矩阵，
        # 我们可以把它看作是初始化了两次，而每次初始化的变换矩阵会做不同的处理，第一次初始化的变换矩阵分布在正弦波上，第二次初始化的变换矩阵分布在余弦波上
        # 并把这两个矩阵分别填充在位置编码矩阵的偶数和奇数位置上，组成最终的位置编码矩阵，

        div_term = torch.exp(torch.arange(0, d_model, 2) * # 2是跳跃，512变256
                             -(math.log(10000.0) / d_model)) # 乘以的是一个缩放的数字

        # 将“变化矩阵”进行奇数和偶数的分别赋值
        pe[:,0::2] = torch.sin(position * div_term) # 所有行，偶数列，用正弦波赋值
        pe[:,1::2] = torch.cos(position * div_term) # 所有行，奇数列，用余弦波赋值

        # 得到编码矩阵pe， 目前pe仍然是二维矩阵，需要和embedding（三维张量）相加
        # 将pe从二维变成三维矩阵
        pe = pe.unsqueeze(0) # 0-行维度增加维度， 1-列维度增加维度

        # 将pe注册成模型的buffer
        # buffer：对模型有帮助的，但不是模型超参数或者参数，也就是不随着训练而被优化的内容
        # 注册后可以随模型一起被保存和重新加载

        self.register_buffer('pe', pe)

    def forward(self, x):
        """
        :param x: 文本序列的词嵌入表示
        :return:
        """
        # pe的长度太长，要进行限制
        # 将3维张量的第2维，缩小成句子的长度，
        # x就是输入的文本，因此pe长度调整为x的第二维度（文本长度），即 x.size(1)
        # max_len == 5000 很大，很难有一个句子长达5000个词汇，pe要与输入张量的适配，将pe做截取（只保留x.size(1)列即可）
        # pe不参与梯度下降训练，因此requires_grad=False

        # x（输入文本） + pe（位置编码）
        x = x + Variable(self.pe[:, :x.size(1)], requires_grad=False)

        # 最后经过 dropout 层
        return self.dropout(x)


def test_dropout():
    m = nn.Dropout(p= 0.2)
    input = torch.randn(4, 5)
    print(input)
    output = m(input)
    print(output)
    # ?? 失活前后矩阵不一致？？


def test_unsqueeze():
    x = torch.tensor([1,2,3,4])
    x1 = torch.unsqueeze(x, 0)
    print(x1)
    x2 = torch.unsqueeze(x, 1)
    print(x2)


def test():
    # 创建 15 X 5 的画布
    plt.figure(figsize=(15, 5))

    x_1 = Variable(torch.LongTensor([[0,1,2,3],[6,7,8,9]]))
    print(x_1)

    # 1个句子，100个词，20维度
    x_2 = Variable(torch.zeros(1, 100, 20))
    print(x_2)

    # 实例化, 20维度，置0比例为0
    pe = PositionalEncoding(20, 0)

    # 因为Input是0矩阵，所以PE只“包含”了位置位置矩阵
    position_matrix = pe(x_2)
    print(position_matrix)

    # 第1位：只有一个元素（句子），
    # 第2位：100个词，全部显示他们的位置信息（PE矩阵体现的就是这100个词的位置信息）
    # 第3位：观察 4 5 6 7 维度 体现出的位置，4和6偶数位应该是Sin，5和7奇数位置应该是Cos
    print(position_matrix[0, :, 4:8])

    # 同一个词，在句子中位置不同，而其位置向量不同
    plt.plot(np.arange(100), position_matrix[0, :, 4:8])
    plt.legend(["dim %d" %p for p in [4,5,6,7]])
    plt.show()
    aa = 0

if __name__ == '__main__':

    test()
    # test_dropout()
    # test_unsqueeze()

